from openai import OpenAI
import json

# 配置 OpenAI 客户端
API_KEY = "sk-zO8exlBicZh7nJeZn5GuC5X9SPuVrZzXoGyOW0i9BFvN62ON" 
BASE_URL = "https://api.chatfire.cn/v1"  

client = OpenAI(
    api_key=API_KEY,
    base_url=BASE_URL
)

def llm_invoke(messages, max_tokens=1500, temperature=0.1):
    """统一的LLM调用接口"""
    try:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature
        )
        return response.choices[0].message.content
    except Exception as e:
        return json.dumps({"error": f"LLM调用失败：{str(e)}"})

# ========== Agent 1：需求分析 ==========
def req_analyzer(user_prompt: str) -> str:
    messages = [
        {
            "role": "system",
            "content": """
            输出格式：
            {
              "functional": [{"name": "功能名称", "description": "功能描述"}],
              "non_functional": [{"name": "非功能名称", "description": "非功能描述"}]
            }
            """
        },
        {
            "role": "user",
            "content": f"分析需求：{user_prompt}"
        }
    ]
    return llm_invoke(messages)

# ========== Agent 2：用例图生成 ==========
def usecase_generator(requirements: str) -> str:
    messages = [
        {
            "role": "system",
            "content": """
            输出格式：
            {
                "system_name": "系统名称",
                "actors": [{"name": "参与者", "description": "描述"}],
                "usecases": [{"name": "用例", "description": "描述", "actors": [], "includes": [], "extends": []}],
                "relationships": [{"from": "A", "to": "B", "type": "association/include/extend"}]
            }
            """
        },
        {
            "role": "user",
            "content": f"根据需求生成用例图：{requirements}"
        }
    ]
    return llm_invoke(messages)

# ========== Agent 3：类图生成 ==========
def class_generator(prompt: str) -> str:
    messages = [
        {
            "role": "system",
            "content": """
            输出格式：
            {
                "classes": [{"name": "类名", "attributes": [{"name": "属性", "type": "类型"}], "methods": [{"name": "方法", "params": [], "return": "返回类型"}]}},
                "associations": [{"from": "类A", "to": "类B", "type": "hasOne/hasMany", "description": "关系描述"}]
            }
            """
        },
        {
            "role": "user",
            "content": prompt
        }
    ]
    return llm_invoke(messages)

# ========== Agent 4：系统顺序图生成 ==========
def sequence_generator(usecase_dsl: str, class_dsl: str) -> str:
    messages = [
        {
            "role": "system",
            "content": """
            输出格式：
            {
                "participants": [{"name": "参与者/类名", "type": "Actor/Class"}],
                "messages": [
                    {
                        "from": "发送者",
                        "to": "接收者",
                        "content": "消息内容",
                        "order": 1
                    }
                ]
            }
            """
        },
        {
            "role": "user",
            "content": f"""
            根据以下内容生成系统顺序图：
            用例图：{usecase_dsl}
            类图：{class_dsl}
            
            要求：
            1. 参与者包括用例的Actor和类图的核心类
            2. 消息序列对应主要用例的执行流程
            3. 消息内容需体现方法调用（如 User.login()）
            """
        }
    ]
    return llm_invoke(messages)

# ========== Agent 5：OCL合约生成 ==========
def ocl_generator(class_dsl: str) -> str:
    messages = [
        {
            "role": "system",
            "content": """
            输出格式：
            {
                "classes": [
                    {
                        "name": "类名",
                        "methods": [
                            {
                                "name": "方法名",
                                "precondition": "OCL前置条件",
                                "postcondition": "OCL后置条件"
                            }
                        ]
                    }
                ]
            }
            """
        },
        {
            "role": "user",
            "content": f"为以下类图生成OCL合约：{class_dsl}"
        }
    ]
    return llm_invoke(messages)


